# -*- coding: utf-8 -*-
# !/usr/bin/python3

import re
import hanlp
import json
from models.Model import Model
from sqlalchemy import and_, or_, not_
from models.HandlePdfPartition import HandlePdfPartitionModel as HandlePdfPartition

# model
model = Model()

# 加载Hanlp模型
tok = hanlp.load(hanlp.pretrained.tok.COARSE_ELECTRA_SMALL_ZH)
ner = hanlp.load(hanlp.pretrained.ner.MSRA_NER_ELECTRA_SMALL_ZH)

# 特殊字符
regex_special = re.compile(r"[…@_!#$%^&*()<>?/\|}{~:]")
# 标点字符
regex_mark = re.compile(r"[^\w\s]")

# 查询条件
condition = and_(
    HandlePdfPartition.status == HandlePdfPartition.status_unprocessed
)

resources = model.all_limit(HandlePdfPartition, condition, 200)
for i in resources:
    print('处理ID：' + str(i.id) + '---------------------')
    personal_name = []
    lines = i.content.splitlines()
    for index, line in enumerate(lines):
        # 为空则跳过
        if len(line) == 0:
            continue
        # 分词
        text_tok = tok(line)
        # 提取人名
        text_ner = ner(text_tok)
        # 收集全部姓名
        for p in text_ner:
            tag = p[1]
            text = p[0]
            if tag == 'PERSON' \
                    and len(text) > 1 \
                    and regex_special.search(text) is None \
                    and regex_mark.search(text) is None:
                personal_name.append(text)
    personal_name = list(set(personal_name))
    i.status = HandlePdfPartition.status_processed
    i.personal_name = json.dumps(personal_name, ensure_ascii=False)
    model.update(i)
del resources
